# $Id: crypto.py 9227e81e919a 2012/08/27 15:49:57 pgurumur $
# Copyright (c) 2009 Prabhu Gurumurthy  <pgurumur@gmail.com>
# 
# Permission to use, copy, modify, and distribute this software for any
# purpose with or without fee is hereby granted, provided that the above
# copyright notice and this permission notice appear in all copies.
#
# THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES
# WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF
# MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR
# ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES
# WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN
# ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF
# OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE.
#

__version__  = "$Revision: 9227e81e919a $"
__author__   = "$Author: pgurumur $"
__modified__ = "$Date: 2012-08-27 15:49:57Z $"

import hashlib
import base64
import string
import os

from Crypto.Cipher import AES, Blowfish, DES3

import stdlib
from exception import CryptoError

CBC = 1
ECB = 2

class crypto(object):
   def __init__(self, Key):
      super(crypto, self).__init__()
      self._err = CryptoError()
      self._key = None
      self._crypt = None

      if Key:
         self._key = hashlib.md5(Key).hexdigest()
      else:
         raise CryptoError("no key defined!")

   def encrypt(self, inputString):
      raise NotImplementedError()

   def decrypt(self, inputString):
      raise NotImplementedError()

class aes(crypto):
   BLOCK_SIZE = AES.block_size
   def __init__(self, Key, Mode = ECB):
      super(aes, self).__init__(Key)
      self._alg = None
      if Mode == CBC:
         self._alg = AES.MODE_CBC
      elif Mode == ECB:
         self._alg = AES.MODE_ECB

      self._iv = "0123456789ABCDEF"
      if self._alg:
         self._crypt = AES.new(self._key, self._alg, self._iv)

      self._eod = hashlib.md5(self._key).hexdigest()

   def _numOfPadBytes(self, BlockSize, Size):
      if not (0 < BlockSize < 255):
         raise CryptoError("block size must be between 0 and 255")

      return BlockSize - (Size % BlockSize)

   def _appendPadding(self, BlockSize, String):
      ix = self._numOfPadBytes(BlockSize, len(String))
      return String + (chr(ix) * ix)

   def _removePadding(self, BlockSize, String):
      ix = ord(String[-1])
      if (ix > BlockSize) or (ix > len(String)):
         raise CryptoError("Invalid padding!")

      return String[:-ix]

   def encrypt(self, inputString):
      encString = None
      if inputString:
         if self._alg == AES.MODE_ECB:
            saltLength = 0
            dataLength = len(inputString) + len(self._eod)
            if dataLength < 16:
               saltLength = 16 - dataLength
            else:
               saltLength = 16 - dataLength % 16

            encString = ''.join([inputString, self._eod,
               stdlib.GetRandomString(saltLength)])
         elif self._alg == AES.MODE_CBC:
            encString = self._appendPadding(self.BLOCK_SIZE, inputString)
      else:
         raise CryptoError("no string given to encrypt")

      # return self._crypt.encrypt(encString).encode("base64")
      return base64.b64encode(self._crypt.encrypt(encString))

   def decrypt(self, inputString):
      retval = None
      if inputString:
         if self._alg == AES.MODE_ECB:
            decString = self._crypt.decrypt(base64.b64decode(inputString))
            retval = decString.split(self._eod)[0]
         elif self._alg == AES.MODE_CBC:
            decString = self._crypt.decrypt(inputString)
            retval = self._removePadding(self.BLOCK_SIZE, decString)

      return retval

class blowfish(crypto):
   def __init__(self, Key):
      super(blowfish, self).__init__(Key)
      self._strLength = 8
      self._crypt = Blowfish.new(self._key)

   def encrypt(self, inputString):
      retval = None
      if inputString:
         repeat = self._strLength - len(inputString) % 8
         inputString = inputString + " " * repeat
         retval = self._crypt.encrypt(inputString)
         
      return retval

   def decrypt(self, inputString):
      retval = None
      if inputString:
         retval = self._crypt.decrypt(inputString).strip()

      return retval

class tripledes(crypto):
   def __init__(self, Key):
      super(tripledes, self).__init__(Key)
      self._strLength = 8
      self._crypt = DES3.new(self._key[:24])

   def encrypt(self, inputString):
      retval = None
      if inputString:
         repeat = self._strLength - len(inputString) % 8
         inputString = inputString + " " * repeat
         retval = self._crypt.encrypt(inputString)
         
      return retval

   def decrypt(self, inputString):
      retval = None
      if inputString:
         retval = self._crypt.decrypt(inputString).strip()

      return retval
